﻿namespace Microsoft.VisualStudio.TestTools.UnitTesting
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Diagnostics.Contracts;
    using System.Linq;
    using System.Threading;
    using System.Windows.Forms;

    /// <summary>
    /// A helper utility for running code on a UI thread.
    /// </summary>
    [DebuggerStepThrough]
    public static class UIHelper
    {
        private static readonly object lockObject = new object();
        private static List<Form> allowedFormsAcrossCalls = new List<Form>();
        private static Form form;
        private static bool started;
        private static EventWaitHandle startWaitHandle;

        #region Public Properties
        /// <summary>
        /// Gets the allowed forms across calls.
        /// </summary>
        /// <value>The allowed forms across calls.</value>
        /// <remarks>
        /// Once the action is run, if there are any forms still alive it is assumed the action failed and the UI
        /// thread needs to be released. If you want a form to survive over multiple actions then add it to the collection.
        /// </remarks>
        public static IList<Form> AllowedFormsAcrossCalls
        {
            get { return UIHelper.allowedFormsAcrossCalls; }
        }
        #endregion Public Properties

        #region Public Methods
        /// <summary>
        /// Releases this UI thread.
        /// </summary>
        public static void Release()
        {
            if (UIHelper.form != null)
            {
                UIHelper.form.Invoke(new Action(UIHelper.form.Close));
                UIHelper.form = null;

                UIHelper.started = false;
            }
        }

        /// <summary>
        /// Runs the action on a UI thread.
        /// </summary>
        /// <param name="action">The action to run.</param>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "We need to catch all exceptions here since they are rethrown.")]
        public static void RunOnUIThread(Action action)
        {
            Contract.Requires(action != null);

            // If we are already on a UI thread then just run the action.
            if (Application.MessageLoop)
            {
                action();
                return;
            }

            EnsureUIThreadIsStarted();

            Exception exception = null;
            UIHelper.form.Invoke(new Action(() =>
            {
                try
                {
                    action();
                }
                catch (Exception ex)
                {
                    exception = ex;
                }
            }));

            // There should only be the test form open after running the action.
            if (UIHelper.NonAllowedFormsStillOpen().Any())
            {
                try
                {
                    // Clean up, allowing the next test to start a new clean UI thread.
                    Release();
                }
                catch
                {
                    if (exception == null)
                    {
                        throw new InvalidOperationException("A form was still open after the action finished.");
                    }
                }
            }

            if (exception != null)
            {
                exception.PreserveStackTrace();
                throw exception;
            }
        }
        #endregion Public Methods

        #region Private Methods
        private static void CurrentDomain_DomainUnload(object sender, EventArgs e)
        {
            Application.Exit();
        }

        private static void EnsureUIThreadIsStarted()
        {
            // If we haven't already started the UI thread then do it now.
            if (!UIHelper.started)
            {
                lock (UIHelper.lockObject)
                {
                    if (!UIHelper.started)
                    {
                        StartUIThread();
                    }
                }
            }
        }

        private static IEnumerable<Form> NonAllowedFormsStillOpen()
        {
            return
                from Form form in Application.OpenForms
                where !object.ReferenceEquals(form, UIHelper.form)
                    && !UIHelper.AllowedFormsAcrossCalls.Contains(form)
                select form;
        }

        private static void Started(object sender, EventArgs e)
        {
            Application.Idle -= new EventHandler(Started);
            UIHelper.started = true;
            UIHelper.startWaitHandle.Set();
        }

        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "The Form stays around to be used by other unit tests.")]
        private static void StartMessageLoop()
        {
            // Create the form we are going to use to invoke the tests.
            UIHelper.form = new Form()
            {
                Text = Properties.Resources.UITestThread
            };

            AppDomain.CurrentDomain.DomainUnload += new EventHandler(CurrentDomain_DomainUnload);

            Application.EnableVisualStyles();
            Application.Idle += new EventHandler(Started);
            Application.Run(UIHelper.form);
        }

        private static void StartUIThread()
        {
            // Create the wait handle that will notify us when the message loop is started.
            UIHelper.startWaitHandle = new EventWaitHandle(false, EventResetMode.ManualReset);

            // Starts the STA UI thread.
            Thread thread = new Thread(new ThreadStart(StartMessageLoop));
            thread.SetApartmentState(ApartmentState.STA);
            thread.Start();

            // Wait for the message loop to start.
            UIHelper.startWaitHandle.WaitOne();
            UIHelper.startWaitHandle.Close();
            UIHelper.startWaitHandle = null;
        }
        #endregion Private Methods
    }
}